import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (classification_report, accuracy_score, 
                            confusion_matrix, roc_curve, auc)
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline

# 配置参数
MAX_FRAMES = 2000  # 每个样本取前2000帧数据
TEST_SIZE = 0.2    # 测试集比例20%
RANDOM_STATE = 42  # 随机种子，保证结果可复现

# 设置中文字体和绘图风格
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
sns.set_style('whitegrid')  # 使用seaborn的白色网格风格

def 加载数据(asd文件夹='ASD', td文件夹='TD'):
    """从ASD和TD文件夹加载数据"""
    特征集 = []
    标签集 = []
    
    # 加载ASD样本
    print("正在加载ASD样本...")
    for 文件 in os.listdir(asd文件夹):
        if 文件.endswith('.csv'):
            数据帧 = pd.read_csv(os.path.join(asd文件夹, 文件))
            # 取前MAX_FRAMES帧
            数据帧 = 数据帧.head(MAX_FRAMES)
            # 提取特征(视线和表情的均值和标准差)
            特征 = [
                数据帧['Gaze_X'].mean(), 数据帧['Gaze_X'].std(),
                数据帧['Gaze_Y'].mean(), 数据帧['Gaze_Y'].std(),
                数据帧['Expression'].mean(), 数据帧['Expression'].std()
            ]
            特征集.append(特征)
            标签集.append(1)  # ASD标签为1
    
    # 加载TD样本
    print("正在加载TD样本...")
    for 文件 in os.listdir(td文件夹):
        if 文件.endswith('.csv'):
            数据帧 = pd.read_csv(os.path.join(td文件夹, 文件))
            # 取前MAX_FRAMES帧
            数据帧 = 数据帧.head(MAX_FRAMES)
            # 提取特征
            特征 = [
                数据帧['Gaze_X'].mean(), 数据帧['Gaze_X'].std(),
                数据帧['Gaze_Y'].mean(), 数据帧['Gaze_Y'].std(),
                数据帧['Expression'].mean(), 数据帧['Expression'].std()
            ]
            特征集.append(特征)
            标签集.append(0)  # TD标签为0
    
    return np.array(特征集), np.array(标签集)

def 绘制特征重要性(特征名称, 重要性值):
    """绘制特征重要性条形图"""
    plt.figure(figsize=(10, 6))
    sns.barplot(x=重要性值, y=特征名称, palette='viridis')
    plt.title('特征重要性分析', fontsize=16, fontweight='bold')
    plt.xlabel('重要性得分', fontsize=12, fontweight='bold')
    plt.ylabel('特征名称', fontsize=12, fontweight='bold')
    plt.xticks(fontsize=10)
    plt.yticks(fontsize=10)
    plt.tight_layout()
    plt.savefig('特征重要性.png', dpi=300, bbox_inches='tight')
    plt.show()

def 绘制混淆矩阵(y真实, y预测, 类别名称):
    """绘制混淆矩阵热力图"""
    混淆矩阵值 = confusion_matrix(y真实, y预测)
    plt.figure(figsize=(8, 6))
    sns.heatmap(混淆矩阵值, annot=True, fmt='d', cmap='Blues',
                xticklabels=类别名称, yticklabels=类别名称,
                annot_kws={"fontsize":12, "fontweight":"bold"})
    plt.title('混淆矩阵', fontsize=16, fontweight='bold')
    plt.xlabel('预测标签', fontsize=12, fontweight='bold')
    plt.ylabel('真实标签', fontsize=12, fontweight='bold')
    plt.xticks(fontsize=10)
    plt.yticks(fontsize=10)
    plt.tight_layout()
    plt.savefig('混淆矩阵.png', dpi=300, bbox_inches='tight')
    plt.show()

def 绘制ROC曲线(y真实, y概率):
    """绘制ROC曲线"""
    fpr, tpr, thresholds = roc_curve(y真实, y概率)
    roc_auc = auc(fpr, tpr)
    
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2,
             label=f'ROC曲线 (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('假阳性率', fontsize=12, fontweight='bold')
    plt.ylabel('真阳性率', fontsize=12, fontweight='bold')
    plt.title('受试者工作特征曲线(ROC)', fontsize=16, fontweight='bold')
    plt.legend(loc="lower right", fontsize=10)
    plt.xticks(fontsize=10)
    plt.yticks(fontsize=10)
    plt.tight_layout()
    plt.savefig('ROC曲线.png', dpi=300, bbox_inches='tight')
    plt.show()

def 绘制分类报告(y真实, y预测):
    """绘制分类报告热力图"""
    报告 = classification_report(y真实, y预测, target_names=['TD', 'ASD'], output_dict=True)
    报告数据 = pd.DataFrame(报告).iloc[:-1, :].T
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(报告数据, annot=True, cmap='YlGnBu', fmt='.2f',
                annot_kws={"fontsize":12, "fontweight":"bold"})
    plt.title('分类报告指标', fontsize=16, fontweight='bold')
    plt.xticks(fontsize=10)
    plt.yticks(fontsize=10)
    plt.tight_layout()
    plt.savefig('分类报告.png', dpi=300, bbox_inches='tight')
    plt.show()

def 主函数():
    # 加载数据
    print("开始数据加载...")
    特征, 标签 = 加载数据()
    print(f"共加载{len(特征)}个样本(ASD:{sum(标签)}个, TD:{len(标签)-sum(标签)}个)")
    
    # 分层划分训练集和测试集
    print("\n划分训练集和测试集...")
    X训练, X测试, y训练, y测试 = train_test_split(
        特征, 标签, test_size=TEST_SIZE, stratify=标签, random_state=RANDOM_STATE
    )
    print(f"训练集: {len(X训练)}个样本")
    print(f"测试集: {len(X测试)}个样本")
    
    # 创建并训练模型(标准化+随机森林)
    print("\n训练模型中...")
    模型 = make_pipeline(
        StandardScaler(),
        RandomForestClassifier(n_estimators=100, random_state=RANDOM_STATE)
    )
    模型.fit(X训练, y训练)
    
    # 评估模型
    print("\n模型评估结果:")
    y预测 = 模型.predict(X测试)
    y概率 = 模型.predict_proba(X测试)[:, 1]  # 获取ASD类的概率
    
    print("\n分类报告:")
    print(classification_report(y测试, y预测, target_names=['TD', 'ASD']))
    print(f"准确率: {accuracy_score(y测试, y预测):.2f}")
    
    # 特征重要性分析
    随机森林 = 模型.named_steps['randomforestclassifier']
    特征名称 = ['Gaze_X均值', 'Gaze_X标准差', 'Gaze_Y均值', 'Gaze_Y标准差', 
              'Expression均值', 'Expression标准差']
    
    print("\n特征重要性分析:")
    for 名称, 重要性 in zip(特征名称, 随机森林.feature_importances_):
        print(f"{名称}: {重要性:.3f}")
    
    # 可视化部分
    print("\n生成可视化图表...")
    绘制特征重要性(特征名称, 随机森林.feature_importances_)
    绘制混淆矩阵(y测试, y预测, ['TD', 'ASD'])
    绘制ROC曲线(y测试, y概率)
    绘制分类报告(y测试, y预测)
    
    print("\n所有可视化图表已保存为PNG文件")

if __name__ == "__main__":
    主函数()